package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"runtime"
	"sync"
	"time"
)

var wg sync.WaitGroup
var successNum int
var failedNum int
var successChan chan error

var duration = flag.Int("d", 0, "Duration of request.The unit is seconds.")
var thread = flag.Int("t", 0, "Number of threads.")
var keepAlive = flag.Bool("k", false, "Use keep-alive for http protocol.")
var reqUrl = flag.String("u", "", "The URL you want to test")
var requestFile = flag.String("f", "", "specify the request definition file.")
var generateSample = flag.Bool("g", false, "generate the request definition file template to the current directory.")

var request Request

type Request struct {
	Url      string            `json:"url"`
	FormData map[string]string `json:"formData"`
	Method   string            `json:"method"`
	Header   map[string]string `json:"header"`
}

func (r Request) Valid() (err error) {
	// 检查 url 格式
	if urlErr := checkUrl(r.Url); urlErr != nil {
		return urlErr
	}

	// 检查 method
	if methodErr := checkHttpMethod(r.Method); methodErr != nil {
		return methodErr
	}

	return
}

func checkUrl(urlAddr string) (err error) {
	if urlAddr == "" {
		return errors.New("url is empty")
	}

	urlObj, urlErr := url.Parse(urlAddr)
	if urlErr != nil {
		return errors.New("An error occurred while parsing the url")
	}

	if urlObj.Scheme != "http" && urlObj.Scheme != "https" {
		return errors.New("url schema illegal:" + urlObj.Scheme)
	}

	//  go run . -d 5 -t 100  -u "http://www.baidu.com:9090/a/b?action=1&test=2"

	//www.baidu.com
	hostName := urlObj.Hostname()
	// 80
	port := urlObj.Port()

	portErr := checkPort(fmt.Sprintf("%s:%s", hostName, port))
	if portErr != nil {
		return portErr
	}

	return
}

func checkPort(ipPorts string) (err error) {
	netRes, netErr := net.DialTimeout("tcp", ipPorts, time.Second*3)
	if netErr != nil {
		return netErr
	}
	if netRes == nil {
		return errors.New("the %s is disabled.")
	}

	if netRes != nil {
		netRes.Close()
		return nil
	}

	return
}

func main() {
	flag.Parse()
	//raising the limits. Some performance gains were achieved with the + goroutines (not a lot).
	runtime.GOMAXPROCS(runtime.NumCPU() + *thread)

	if *generateSample {
		err := GenerateRequestFile("./request_sample.json")
		if err != nil {
			fmt.Println(err)
		}
		return
	}

	// 检查参数并使用 flag.Parse 解析命令行输入
	request := checkParams()

	validErr := request.Valid()
	if validErr != nil {
		log.Fatal(validErr)
	}

	fmt.Printf("%+v \n", request)

	// 使用该通道来存储请求的结果,并启用一个协程来读取该通道的结果
	successChan = make(chan error, *thread)
	counterCtx, counterCancel := context.WithCancel(context.Background())
	defer counterCancel()
	go counter(counterCtx)

	start := time.Now()

	for i := 1; i <= *thread; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx := context.TODO()
			ctx, cancel := context.WithTimeout(ctx, time.Duration(1e9*(*duration)))
			defer cancel()
		LOOP:
			for {
				httpCtx := context.TODO()
				client := getHttpClient(*keepAlive, time.Duration(1e9*(*duration)))
				_, err := HandleReq(httpCtx, client, request)

				// 保存结果
				successChan <- err

				select {
				case <-ctx.Done():
					// 结束执行
					break LOOP
				default:
					// default 使select 避免陷入阻塞
				}
			}

		}()
	}
	wg.Wait()
	elapsed := time.Since(start)

	perSecondTimes := float64(successNum) / elapsed.Seconds()
	fmt.Printf("%g / %g = %g \n", float64(successNum), elapsed.Seconds(), perSecondTimes)
	fmt.Printf("程序执行了%g秒,req/s = %g,success:%d, failed:%d \n",
		elapsed.Seconds(),
		perSecondTimes,
		successNum,
		failedNum,
	)

}

func checkParams() (request Request) {
	if *duration == 0 || *thread == 0 {
		log.Fatal("params is error.Use \"-d\" and \"-t\" parameter add the set.")
	}
	if *requestFile == "" && *reqUrl == "" {
		log.Fatal("the URL cannot be empty.Use the \"-u\" or \"-f\" parameter to set the URL.")
	}

	if *requestFile != "" && *reqUrl != "" {
		log.Fatal("the \"-u\" or \"-f\" parameter can not exist the same time.")
	}

	if *requestFile != "" {
		fileBytes, err := ioutil.ReadFile(*requestFile)
		if err != nil {
			log.Fatal("an error occurred reading the file", err)
		}
		unmarshalErr := json.Unmarshal(fileBytes, &request)
		if unmarshalErr != nil {
			log.Fatal("unmarshal err: ", unmarshalErr)
		}
	}

	if *reqUrl != "" {
		request.Url = *reqUrl
		request.Method = http.MethodGet
	}

	return
}

func counter(ctx context.Context) {
LOOP:
	for {
		select {
		case <-ctx.Done():
			fmt.Printf("成功的通道里的剩余数量:len = %d \n", len(successChan))
			break LOOP
		case err := <-successChan:
			if err == nil {
				successNum += 1
			} else {
				failedNum += 1
				fmt.Printf("request err:%+v \n", err)
			}
		default:
		}
	}
}
